# plotly_trained_states_3d_with_best.py

import json
import plotly.graph_objs as go
import os
import numpy as np

def load_trained_states(trained_states_file='trained_states.json'):
    """
    加载训练过程中访问过的状态及其对应的值。

    参数:
    - trained_states_file (str): 存储训练状态的JSON文件路径。

    返回:
    - trained_states (list of dict): 每个元素包含超参数组合和对应的准确率。
    """
    if os.path.exists(trained_states_file):
        with open(trained_states_file, 'r') as f:
            trained_states = json.load(f)
        print(f"Loaded {len(trained_states)} trained states from {trained_states_file}.")
        return trained_states
    else:
        print(f"No trained states found at {trained_states_file}.")
        return []

def load_human_optimal_states(human_optimal_file='human_optimal_states.json'):
    """
    加载人类定义的最优状态及其对应的值。

    参数:
    - human_optimal_file (str): 存储人类最优状态的JSON文件路径。

    返回:
    - human_optimal_states (list of dict): 每个元素包含超参数组合和对应的准确率。
    """
    if os.path.exists(human_optimal_file):
        with open(human_optimal_file, 'r') as f:
            human_optimal_states = json.load(f)
        print(f"Loaded {len(human_optimal_states)} human optimal states from {human_optimal_file}.")
        return human_optimal_states
    else:
        print(f"No human optimal states found at {human_optimal_file}.")
        return []

def plot_interactive_trained_states_3d_with_best(trained_states, human_optimal_states):
    """
    绘制交互式3D散点图，突出显示RL和人类的最佳状态。

    参数:
    - trained_states (list of dict): 包含超参数组合和对应准确率的训练状态。
    - human_optimal_states (list of dict): 包含超参数组合和对应准确率的人类最优状态。

    返回:
    - None
    """
    if not trained_states:
        print("No trained states data available to plot.")
        return

    # 提取训练状态数据
    lr_trained = [state['learning_rate'] for state in trained_states]
    md_trained = [state['max_depth'] for state in trained_states]
    ne_trained = [state['n_estimators'] for state in trained_states]
    acc_trained = [state['accuracy'] for state in trained_states]

    lr_trained = np.array(lr_trained)
    md_trained = np.array(md_trained)
    ne_trained = np.array(ne_trained)
    acc_trained = np.array(acc_trained)

    # 识别RL最佳状态
    max_accuracy_rl = np.max(acc_trained)
    best_indices_rl = np.where(acc_trained == max_accuracy_rl)[0]
    best_lr_rl = lr_trained[best_indices_rl]
    best_md_rl = md_trained[best_indices_rl]
    best_ne_rl = ne_trained[best_indices_rl]
    best_acc_rl = acc_trained[best_indices_rl]

    # 加载并提取人类最优状态数据
    if human_optimal_states:
        lr_human = [state['learning_rate'] for state in human_optimal_states]
        md_human = [state['max_depth'] for state in human_optimal_states]
        ne_human = [state['n_estimators'] for state in human_optimal_states]
        acc_human = [state['accuracy'] for state in human_optimal_states]

        lr_human = np.array(lr_human)
        md_human = np.array(md_human)
        ne_human = np.array(ne_human)
        acc_human = np.array(acc_human)

        # 识别人类最佳状态
        max_accuracy_human = np.max(acc_human)
        best_indices_human = np.where(acc_human == max_accuracy_human)[0]
        best_lr_human = lr_human[best_indices_human]
        best_md_human = md_human[best_indices_human]
        best_ne_human = ne_human[best_indices_human]
        best_acc_human = acc_human[best_indices_human]
    else:
        best_lr_human, best_md_human, best_ne_human, best_acc_human = [], [], [], []

    # 创建训练状态的散点图
    scatter_trained = go.Scatter3d(
        x=lr_trained,
        y=md_trained,
        z=ne_trained,
        mode='markers',
        marker=dict(
            size=5,
            color=acc_trained,
            colorscale='Viridis',
            opacity=0.8,
            colorbar=dict(title='Validation Accuracy')
        ),
        text=[f'RL Visited State<br>LR: {lr}<br>MD: {md}<br>NE: {ne}<br>Acc: {acc:.4f}'
              for lr, md, ne, acc in zip(lr_trained, md_trained, ne_trained, acc_trained)],
        hoverinfo='text',
        name='Visited States'
    )

    # 创建RL最佳状态的散点图
    best_scatter_rl = go.Scatter3d(
        x=best_lr_rl,
        y=best_md_rl,
        z=best_ne_rl,
        mode='markers',
        marker=dict(
            size=10,
            color='red',
            symbol='diamond',  # 使用有效符号
            opacity=1.0,
            line=dict(
                color='black',
                width=2
            )
        ),
        text=[f'RL Best State<br>LR: {lr}<br>MD: {md}<br>NE: {ne}<br>Acc: {acc:.4f}'
              for lr, md, ne, acc in zip(best_lr_rl, best_md_rl, best_ne_rl, best_acc_rl)],
        hoverinfo='text',
        name='RL Best States'
    )

    # 创建人类最佳状态的散点图
    if len(best_lr_human) > 0:
        best_scatter_human = go.Scatter3d(
            x=best_lr_human,
            y=best_md_human,
            z=best_ne_human,
            mode='markers',
            marker=dict(
                size=10,
                color='blue',
                symbol='x',  # 使用不同的有效符号
                opacity=1.0,
                line=dict(
                    color='black',
                    width=2
                )
            ),
            text=[f'Human Best State<br>LR: {lr}<br>MD: {md}<br>NE: {ne}<br>Acc: {acc:.4f}'
                  for lr, md, ne, acc in zip(best_lr_human, best_md_human, best_ne_human, best_acc_human)],
            hoverinfo='text',
            name='Human Best States'
        )
    else:
        best_scatter_human = None

    # 合并所有散点图
    data = [scatter_trained, best_scatter_rl]
    if best_scatter_human:
        data.append(best_scatter_human)

    # 创建布局
    layout = go.Layout(
        title='Interactive 3D Scatter Plot: Trained States with Best States Highlighted',
        scene=dict(
            xaxis=dict(title='Learning Rate'),
            yaxis=dict(title='Max Depth'),
            zaxis=dict(title='Number of Estimators')
        ),
        legend=dict(x=0, y=1)
    )

    fig = go.Figure(data=data, layout=layout)
    fig.show()

if __name__ == "__main__":
    trained_states = load_trained_states('trained_states.json')
    human_optimal_states = load_human_optimal_states('human_optimal_states.json')
    plot_interactive_trained_states_3d_with_best(trained_states, human_optimal_states)
